# -*- coding: utf-8 -*-
#
# Copyright (c) 2025 Huawei Technologies Co., Ltd.
# openUBMC is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#         http://license.coscl.org.cn/MulanPSL2
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.

"""
操作符重载发射器

处理需要操作符重载的比较操作，如 Optional[T] != None -> y.has_value()
"""

from typing import Optional
from ...ir import IRNode
from ...ir.expressions import IRCompare, IRBinOp, IRUnaryOp, IRConstant
from ...ir.expressions.literals import IRList, IRDict, IRSet, IRTuple
from ..core.emitter import BaseEmitter, EmitResult
from ..core.context import CodeGenContext
from ..configs.operators import get_operator_precedence
from ..core.generator import get_global_generator
from ...symbol_meta import TypeInfo, TypeKind
from ...template_system import replace_template_params


class OperatorOverloadEmitter(BaseEmitter):
    """操作符重载发射器

    专门处理需要操作符重载的情况，如：
    - Optional[T] != None -> obj.has_value()
    - Optional[T] == None -> !obj.has_value()
    - not list -> !list.empty()
    """

    # 临时变量计数器（类级别，与 TemplateEmitter 共享）
    _temp_var_counter = 0

    @staticmethod  # type: ignore
    def _count_placeholder_occurrences(template: str, placeholder: str) -> int:
        """统计模板中占位符出现的次数"""
        return template.count(placeholder)

    @staticmethod  # type: ignore
    def _is_temporary_container(ir_node: IRNode) -> bool:
        """判断IR节点是否是临时容器字面值"""
        return isinstance(ir_node, (IRList, IRDict, IRSet, IRTuple))

    @classmethod  # type: ignore
    def _get_next_temp_var_name(cls) -> str:
        """生成唯一的临时变量名"""
        cls._temp_var_counter += 1
        return f"__mcpy_temp_{cls._temp_var_counter}"

    def can_emit(self, node: IRNode, context: CodeGenContext) -> bool:
        """检查是否为需要操作符重载的操作"""
        if isinstance(node, IRCompare):
            return self._can_emit_compare(node, context)
        elif isinstance(node, IRBinOp):
            return self._can_emit_binop(node, context)
        elif isinstance(node, IRUnaryOp):
            return self._can_emit_unaryop(node, context)
        else:
            return False

    def _can_emit_compare(self, node: IRCompare, context: CodeGenContext) -> bool:
        """检查是否可以处理比较操作"""
        # 只处理单个比较操作（不处理链式比较）
        if len(node.ops) != 1 or len(node.comparators) != 1:
            return False

        op = node.ops[0]
        # 处理所有比较操作符
        if op not in ["==", "!=", "is", "is not", "in", "not in", "<", "<=", ">", ">="]:
            return False

        # 不在 can_emit 中执行符号查找，太昂贵
        # 直接返回 True，让 emit 方法尝试处理
        # 如果处理失败，emit 会返回 None，发射器系统会继续尝试其他发射器
        return True

    def _can_emit_binop(self, node: IRBinOp, context: CodeGenContext) -> bool:
        """检查是否可以处理二元操作"""
        # 支持的二元操作符
        supported_ops = [
            # 算术操作符
            "+",
            "-",
            "*",
            "/",
            "//",
            "%",
            "**",
            # 位操作符
            "|",
            "&",
            "^",
            "<<",
            ">>",
        ]
        if node.op not in supported_ops:
            return False

        # 不在 can_emit 中执行符号查找，太昂贵
        # 直接返回 True，让 emit 方法尝试处理
        return True

    def _can_emit_unaryop(self, node: IRUnaryOp, context: CodeGenContext) -> bool:
        """检查是否可以处理一元操作"""
        # 只处理 not 操作符
        if node.op != "not":
            return False

        # 检查是否有 __bool__ 方法重载
        return self._has_unaryop_overload(node, context)

    def emit(self, node: IRNode, context: CodeGenContext) -> EmitResult | None:
        """生成操作符重载代码

        Returns:
            EmitResult: 如果找到重载则返回结果
            None: 如果没有找到重载，让其他发射器处理
        """
        if isinstance(node, IRCompare):
            result = self._emit_compare(node, context)
            return result if result else None
        elif isinstance(node, IRBinOp):
            result = self._emit_binop(node, context)
            return result if result else None
        elif isinstance(node, IRUnaryOp):
            return self._emit_unaryop(node, context)
        else:
            # 不支持的节点类型，返回 None
            return None

    def _emit_compare(
        self, node: IRCompare, context: CodeGenContext
    ) -> EmitResult | None:
        """生成比较操作符重载代码"""
        op = node.ops[0]
        left = node.left
        right = node.comparators[0]

        # 尝试查找操作符重载
        overload_result = self._try_compare_overload(left, op, right, context)
        if overload_result:
            return overload_result

        # 如果没有找到重载，返回 None 让其他发射器处理
        return None

    def _emit_binop(self, node: IRBinOp, context: CodeGenContext) -> EmitResult | None:
        """生成二元操作符重载代码"""
        op = node.op
        left = node.left
        right = node.right

        # 尝试查找操作符重载
        overload_result = self._try_binop_overload(left, op, right, context)
        if overload_result:
            return overload_result

        # 如果没有找到重载，返回 None 让其他发射器处理
        return None

    def _emit_unaryop(self, node: IRUnaryOp, context: CodeGenContext) -> EmitResult:
        """生成一元操作符重载代码"""
        op = node.op
        operand = node.operand

        # 尝试查找操作符重载
        overload_result = self._try_unaryop_overload(operand, op, context)
        if overload_result:
            return overload_result

        # 如果没有找到重载，返回默认处理
        return EmitResult(
            code="/* No unaryop overload found */",
            precedence=get_operator_precedence("unary"),
        )

    def _has_compare_overload(self, node: IRCompare, context: CodeGenContext) -> bool:
        """检查是否存在比较操作符重载"""
        op = node.ops[0]
        left = node.left
        right = node.comparators[0]

        return self._try_compare_overload(left, op, right, context) is not None

    def _has_binop_overload(self, node: IRBinOp, context: CodeGenContext) -> bool:
        """检查是否存在二元操作符重载"""
        op = node.op
        left = node.left
        right = node.right

        return self._try_binop_overload(left, op, right, context) is not None

    def _has_unaryop_overload(self, node: IRUnaryOp, context: CodeGenContext) -> bool:
        """检查是否存在一元操作符重载"""
        op = node.op
        operand = node.operand

        return self._try_unaryop_overload(operand, op, context) is not None

    def _try_compare_overload(
        self, left: IRNode, op: str, right: IRNode, context: CodeGenContext
    ) -> Optional[EmitResult]:
        """尝试比较操作符重载"""
        # 处理 in 和 not in 运算符
        if op in ["in", "not in"]:
            return self._handle_membership_test(left, op, right, context)

        # 获取左操作数的类型
        left_type = self._get_node_type(left, context)
        if not left_type:
            return None

        # 检查是否是与 None 的比较
        if self._is_none_constant(right):
            return self._handle_none_comparison(left, left_type, op, context)

        # 处理相等/不等比较操作符 (==, !=)
        if op in ["==", "!="]:
            return self._handle_equality_overload(left, left_type, op, right, context)

        # 处理其他比较操作符 (<, <=, >, >=)
        if op in ["<", "<=", ">", ">="]:
            return self._handle_comparison_overload(left, left_type, op, right, context)

        return None

    def _try_binop_overload(
        self, left: IRNode, op: str, right: IRNode, context: CodeGenContext
    ) -> Optional[EmitResult]:
        """尝试二元操作符重载"""
        # 获取左操作数的类型
        left_type = self._get_node_type(left, context)
        if not left_type:
            return None

        # 查找操作符重载配置
        return self._handle_binop_overload(left, left_type, op, right, context)

    def _try_unaryop_overload(
        self, operand: IRNode, op: str, context: CodeGenContext
    ) -> Optional[EmitResult]:
        """尝试一元操作符重载"""
        # 只处理 not 操作符
        if op != "not":
            return None

        # 获取操作数的类型
        operand_type = self._get_node_type(operand, context)
        if not operand_type:
            return None

        # 查找 __bool__ 方法重载
        return self._handle_bool_overload(operand, operand_type, context)

    def _get_node_type(self, node: IRNode, context: CodeGenContext) -> Optional[str]:
        """获取节点的类型信息"""
        # 尝试从扩展参数获取推导的类型
        type_info = node.get_extension_param("inferred_type")
        if type_info and hasattr(type_info, "name"):
            return str(type_info.name)

        # 如果是名称节点，尝试通过符号表查找
        if hasattr(node, "id") and context.current_node:
            symbol = context.current_node.lookup_symbol(node.id)  # type: ignore
            if symbol and hasattr(symbol, "type_info") and symbol.type_info:
                return str(symbol.type_info.name)

        return None

    def _is_none_constant(self, node: IRNode) -> bool:
        """检查是否是 None 常量"""
        return isinstance(node, IRConstant) and node.value is None

    def _handle_membership_test(
        self, item: IRNode, op: str, container: IRNode, context: CodeGenContext
    ) -> Optional[EmitResult]:
        """处理成员检测运算符 (in / not in)

        item in container -> container.__contains__(item)
        item not in container -> !container.__contains__(item)
        """
        if not context.current_node:
            return None

        # 从推断类型获取容器的 TypeInfo
        container_type_info = container.get_extension_param("inferred_type")

        if not container_type_info:
            # 尝试从容器节点查找符号
            if hasattr(container, "id"):
                container_id = getattr(container, "id")
                symbol = context.current_node.lookup_symbol(container_id)
                if symbol and hasattr(symbol, "type_info"):
                    container_type_info = symbol.type_info

            if not container_type_info:
                return None

        # 获取 item 的类型用于方法签名匹配
        item_type_info = item.get_extension_param("inferred_type")
        args_types = [item_type_info] if item_type_info else []

        # 查找 __contains__ 方法符号
        method_symbol = None
        if context.current_node:
            method_symbol = context.current_node.lookup_symbol(
                "__contains__",
                owner_type=container_type_info,
                func_args=args_types,
            )

        if not method_symbol:
            return None

        # 查找匹配的签名
        matching_signature = method_symbol.find_matching_signature(
            [container_type_info] + args_types
        )
        if not matching_signature:
            # 如果没有找到匹配的签名，使用第一个签名
            if method_symbol.signatures:
                matching_signature = method_symbol.signatures[0]
            else:
                return None

        # 获取模板
        template = matching_signature.emit_template
        if not template:
            return None

        generator = get_global_generator()
        container_code = generator.generate_expression(container, context)
        item_code = generator.generate_expression(item, context)

        # 获取实际的参数名（从签名中提取，第一个参数是 self，第二个是 item/sub/key 等）
        param_name = "item"  # 默认参数名
        if hasattr(matching_signature, "args") and len(matching_signature.args) >= 2:
            # args[0] 是 self，args[1] 是实际参数（可能是 item、sub、key 等）
            param_name = matching_signature.args[1]

        # 兜底机制：检测多个 {self} 展开 + 临时容器字面值的场景
        # 用立即执行lambda包装，避免多次求值导致的悬垂引用
        self_count = self._count_placeholder_occurrences(template, "{self}")
        is_temp_container = self._is_temporary_container(container)

        if self_count >= 2 and is_temp_container:
            # 生成唯一的临时变量名
            temp_var = self._get_next_temp_var_name()

            # 替换模板中的占位符为临时变量名
            template_with_temp = template.replace("{self}", temp_var).replace(
                f"{{{param_name}}}", item_code
            )

            # 用立即执行lambda包装
            code = f"[&]() {{ auto {temp_var} = {container_code}; return {template_with_temp}; }}()"
        else:
            # 正常替换模板占位符
            code = template.replace("{self}", container_code).replace(
                f"{{{param_name}}}", item_code
            )

        # 如果是 not in，在外面包一层 !
        if op == "not in":
            code = f"!({code})"

        # 添加必要的包含文件
        includes = getattr(matching_signature, "includes", [])
        for include in includes:
            context.require_include(include)

        return EmitResult(code=code, precedence=get_operator_precedence("compare"))

    def _handle_none_comparison(
        self, left: IRNode, left_type: str, op: str, context: CodeGenContext
    ) -> Optional[EmitResult]:
        """处理与 None 的比较

        使用符号体系查找，支持 builtin 类型和用户自定义类
        """
        if not context.current_node:
            return None

        # 从推断类型获取 TypeInfo
        left_type_info = left.get_extension_param("inferred_type")
        if not left_type_info:
            return None

        # 根据操作符选择方法
        method_name = None
        if op in ["!=", "is not"]:
            method_name = "__ne__"
        elif op in ["==", "is"]:
            method_name = "__eq__"

        if not method_name:
            return None

        none_type = TypeInfo(kind=TypeKind.NONE, name="None")

        method_symbol = context.current_node.lookup_symbol(
            method_name,
            owner_type=left_type_info,
            func_args=[none_type],
        )

        if not method_symbol:
            return None

        # 查找匹配的签名
        matching_signature = method_symbol.find_matching_signature(
            [left_type_info, none_type]
        )
        if not matching_signature:
            # 如果没有找到匹配的签名，使用第一个签名
            if method_symbol.signatures:
                matching_signature = method_symbol.signatures[0]
            else:
                return None

        # 获取模板
        template = matching_signature.emit_template
        if not template:
            return None

        generator = get_global_generator()
        left_code = generator.generate_expression(left, context)

        # 替换模板中的占位符
        code = template.replace("{self}", left_code)

        # 添加必要的包含文件
        includes = getattr(matching_signature, "includes", [])
        for include in includes:
            context.require_include(include)

        return EmitResult(code=code, precedence=get_operator_precedence("compare"))

    def _handle_binop_overload(
        self,
        left: IRNode,
        left_type: str,
        op: str,
        right: IRNode,
        context: CodeGenContext,
    ) -> Optional[EmitResult]:
        """处理二元操作符重载

        使用符号体系查找，支持 builtin 类型和用户自定义类
        """
        if not context.current_node:
            return None

        # 从推断类型获取 TypeInfo
        left_type_info = left.get_extension_param("inferred_type")
        if not left_type_info:
            return None

        # 操作符到方法名的映射
        op_to_method = {
            # 算术操作符
            "+": "__add__",
            "-": "__sub__",
            "*": "__mul__",
            "/": "__truediv__",
            "//": "__floordiv__",
            "%": "__mod__",
            "**": "__pow__",
            # 位操作符
            "|": "__or__",
            "&": "__and__",
            "^": "__xor__",
            "<<": "__lshift__",
            ">>": "__rshift__",
        }

        method_name = op_to_method.get(op)
        if not method_name:
            return None

        # 获取右操作数的类型用于重载匹配
        right_type_info = right.get_extension_param("inferred_type")
        # args_types 不包含 self（只有方法参数）
        args_types = [right_type_info] if right_type_info else []
        # 完整的 args_types 包含 self（用于签名匹配）
        full_args_types = [left_type_info] + args_types

        # 查找方法符号
        method_symbol = None
        if context.current_node:
            method_symbol = context.current_node.lookup_symbol(
                method_name,
                owner_type=left_type_info,
                func_args=args_types,
            )

        if not method_symbol:
            return None

        # 查找匹配的签名（需要包含 self）
        matching_signature = method_symbol.find_matching_signature(full_args_types)
        if not matching_signature:
            # 如果没有找到匹配的签名，使用第一个签名
            if method_symbol.signatures:
                matching_signature = method_symbol.signatures[0]
            else:
                return None

        # 获取模板
        template = matching_signature.emit_template
        if not template:
            return None

        generator = get_global_generator()
        left_code = generator.generate_expression(left, context)
        right_code = generator.generate_expression(right, context)

        # 替换模板中的占位符
        # 使用签名中的实际参数名进行替换
        code = replace_template_params(
            template, matching_signature, [left_code, right_code]
        )

        # 添加必要的包含文件
        includes = getattr(matching_signature, "includes", [])
        for include in includes:
            context.require_include(include)

        return EmitResult(code=code, precedence=get_operator_precedence(op))

    def _handle_bool_overload(
        self,
        operand: IRNode,
        operand_type: str,
        context: CodeGenContext,
    ) -> Optional[EmitResult]:
        """处理 __bool__ 方法重载

        通过符号体系查找 __bool__ 方法，支持 builtin 类型和用户自定义类
        """
        if not context.current_node:
            return None

        # 从推断类型获取 TypeInfo
        type_info = operand.get_extension_param("inferred_type")
        if not type_info:
            return None

        # 查找 __bool__ 方法符号
        method_symbol = context.current_node.lookup_symbol(
            "__bool__",
            owner_type=type_info,
            func_args=[],
        )

        if not method_symbol:
            return None

        # 获取匹配的签名（现在 args_types 不包含 self）
        matching_signature = method_symbol.find_matching_signature([type_info])
        if not matching_signature:
            # 如果没有找到匹配的签名，使用第一个签名
            if method_symbol.signatures:
                matching_signature = method_symbol.signatures[0]
            else:
                return None

        # 获取模板
        template = matching_signature.emit_template
        if not template:
            return None

        generator = get_global_generator()
        operand_code = generator.generate_expression(operand, context)

        # 替换模板中的占位符
        code = template.replace("{self}", operand_code)

        # 对于 not 操作，需要对 __bool__ 的结果取反
        # 如果模板已经包含 ! 前缀（如 "!{self}.empty()"），则去掉它
        # 否则添加 ! 前缀
        code = code.strip()
        if code.startswith("!"):
            # 去掉前缀的 !，实现取反
            code = code[1:].strip()
        else:
            # 添加 ! 前缀
            code = f"!({code})"

        return EmitResult(code=code, precedence=get_operator_precedence("unary"))

    def _handle_equality_overload(
        self,
        left: IRNode,
        left_type: str,
        op: str,
        right: IRNode,
        context: CodeGenContext,
    ) -> Optional[EmitResult]:
        """处理相等/不等比较操作符重载 (==, !=)

        使用符号体系查找，支持 builtin 类型和用户自定义类
        """
        if not context.current_node:
            return None

        # 从推断类型获取 TypeInfo
        left_type_info = left.get_extension_param("inferred_type")
        if not left_type_info:
            return None

        # 操作符到方法名的映射
        op_to_method = {
            "==": "__eq__",
            "!=": "__ne__",
        }

        method_name = op_to_method.get(op)
        if not method_name:
            return None

        # 获取右操作数的类型用于重载匹配
        right_type_info = right.get_extension_param("inferred_type")
        # args_types 不包含 self（只有方法参数）
        args_types = [right_type_info] if right_type_info else []
        # 完整的 args_types 包含 self（用于签名匹配）
        full_args_types = [left_type_info] + args_types

        # 查找方法符号
        method_symbol = None
        if context.current_node:
            method_symbol = context.current_node.lookup_symbol(
                method_name,
                owner_type=left_type_info,
                func_args=args_types,
            )

        if not method_symbol:
            return None

        # 查找匹配的签名（需要包含 self）
        matching_signature = method_symbol.find_matching_signature(full_args_types)
        if not matching_signature:
            # 如果没有找到匹配的签名，使用第一个签名
            if method_symbol.signatures:
                matching_signature = method_symbol.signatures[0]
            else:
                return None

        # 获取模板
        template = matching_signature.emit_template
        if not template:
            return None

        generator = get_global_generator()
        left_code = generator.generate_expression(left, context)
        right_code = generator.generate_expression(right, context)

        # 替换模板中的占位符
        # 注意：用户自定义类使用双层大括号 {{var}}，内置类型使用单层大括号 {var}
        code = template.replace("{{self}}", left_code).replace("{self}", left_code)
        code = code.replace("{{other}}", right_code).replace("{other}", right_code)

        # 添加必要的包含文件
        includes = getattr(matching_signature, "includes", [])
        for include in includes:
            context.require_include(include)

        return EmitResult(code=code, precedence=get_operator_precedence("compare"))

    def _handle_comparison_overload(
        self,
        left: IRNode,
        left_type: str,
        op: str,
        right: IRNode,
        context: CodeGenContext,
    ) -> Optional[EmitResult]:
        """处理比较操作符重载 (<, <=, >, >=)

        使用符号体系查找，支持 builtin 类型和用户自定义类
        """
        if not context.current_node:
            return None

        # 从推断类型获取 TypeInfo
        left_type_info = left.get_extension_param("inferred_type")
        if not left_type_info:
            return None

        # 操作符到方法名的映射
        op_to_method = {
            "<": "__lt__",
            "<=": "__le__",
            ">": "__gt__",
            ">=": "__ge__",
        }

        method_name = op_to_method.get(op)
        if not method_name:
            return None

        # 获取右操作数的类型用于重载匹配
        right_type_info = right.get_extension_param("inferred_type")
        # args_types 不包含 self（只有方法参数）
        args_types = [right_type_info] if right_type_info else []
        # 完整的 args_types 包含 self（用于签名匹配）
        full_args_types = [left_type_info] + args_types

        # 查找方法符号
        method_symbol = None
        if context.current_node:
            method_symbol = context.current_node.lookup_symbol(
                method_name,
                owner_type=left_type_info,
                func_args=args_types,
            )

        if not method_symbol:
            return None

        # 查找匹配的签名（需要包含 self）
        matching_signature = method_symbol.find_matching_signature(full_args_types)
        if not matching_signature:
            # 如果没有找到匹配的签名，使用第一个签名
            if method_symbol.signatures:
                matching_signature = method_symbol.signatures[0]
            else:
                return None

        # 获取模板
        template = matching_signature.emit_template
        if not template:
            return None

        generator = get_global_generator()
        left_code = generator.generate_expression(left, context)
        right_code = generator.generate_expression(right, context)

        # 替换模板中的占位符
        # 注意：用户自定义类使用双层大括号 {{var}}，内置类型使用单层大括号 {var}
        code = template.replace("{{self}}", left_code).replace("{self}", left_code)
        code = code.replace("{{other}}", right_code).replace("{other}", right_code)

        # 添加必要的包含文件
        includes = getattr(matching_signature, "includes", [])
        for include in includes:
            context.require_include(include)

        return EmitResult(code=code, precedence=get_operator_precedence("compare"))

    def get_priority(self) -> int:
        """获取发射器优先级"""
        return 1500  # 高优先级，在CallbackEmitterAdapter和普通OperatorEmitter之间
